#!/bin/bash
################################################################
##
## Build Release Image 脚本，用于生成指定配置的系统镜像文件
##
################################################################
set -e  # 一旦出现错误则退出脚本

# 如果环境变量ROOT未设置，则将ROOT设置为上级目录的绝对路径
if [ -z $ROOT ]; then
	ROOT=`cd .. && pwd`
fi

# 设置发行版名称，默认为"xineal"，否则使用第一个参数指定的发行版名称
if [ -z $1 ]; then
	DISTRO="xineal"
else
	DISTRO=$1
fi

# 设置平台名称，默认为"rk3399"，否则使用第二个参数指定的平台名称
if [ -z $2 ]; then
	PLATFORM="rk3399"
else
	PLATFORM=$2
fi

# 根据第三个参数决定镜像类型，"1"为桌面版，磁盘大小设为3800MiB，否则为服务器版，磁盘大小1200MiB
if [ $3 = "1" ]; then
	IMAGETYPE="desktop"
	disk_size="3800"
else
	IMAGETYPE="server"
	disk_size="2000"
fi

# 脚本失败自动执行函数
cleanup() {
    local exit_code=$?
    if [ $exit_code -ne 0 ]; then   # 非零表示脚本失败
        echo "Script failed with exit code $exit_code, performing cleanup..."
        if mountpoint -q /media/tmp; then
            umount /media/tmp
            echo "/media/tmp unmounted."
        fi
    fi
}
trap cleanup EXIT

# 定义构建相关路径变量
BUILD="$ROOT/external"
TOOLPATH=${BUILD}/rkbin/tools   # 工具路径
OUTPUT="$ROOT/output"           # 输出路径
VER="v1.0"                     # 版本号

# 镜像名称格式为 OrangePi_${PLATFORM}_${DISTRO}_${IMAGETYPE}_${VER}.img
IMAGENAME="OrangePi_${PLATFORM}_${DISTRO}_${IMAGETYPE}_${VER}.img"
IMAGE="$OUTPUT/images/$IMAGENAME"  # 镜像文件完整路径
ROOTFS="$OUTPUT/rootfs"             # 根文件系统目录

VAR=$4                    # 第四个参数，一般作为后续脚本的参数
PATH=$PATH:$TOOLPATH      # 将工具路径加入PATH环境变量，方便后续调用工具

# 导入分区相关参数（如LOADER1_START，RESERVED1_START等）
source $ROOT/scripts/partitions.sh

# 如果镜像输出目录不存在，则创建
if [ ! -d $OUTPUT/images ]; then
        mkdir -p $OUTPUT/images
fi

# 如果没有设定disk_size，默认为100 MiB
if [ -z "$disk_size" ]; then
	disk_size=100 #MiB
fi

# 磁盘大小须至少为60 MiB，否则脚本报错退出
if [ "$disk_size" -lt 60 ]; then
	echo "Disk size must be at least 60 MiB"
	exit 2
fi

echo "Generate System image : ${IMAGE} !"

# 创建一个空镜像文件，使用dd跳过4000 MB大小（seek=4000），不写入数据（count=0）
dd if=/dev/zero of=${IMAGE} bs=1M count=0 seek=$(( disk_size * 2 ))

# 使用parted创建GPT分区表
parted -s ${IMAGE} mklabel gpt 

# 按分区脚本变量，依次创建分区，单位为扇区(s)
parted -s ${IMAGE} unit s mkpart loader1 ${LOADER1_START} $(expr ${RESERVED1_START} - 1)
parted -s ${IMAGE} unit s mkpart reserved1 ${RESERVED1_START} $(expr ${RESERVED2_START} - 1)
parted -s ${IMAGE} unit s mkpart reserved2 ${RESERVED2_START} $(expr ${LOADER2_START} - 1)
parted -s ${IMAGE} unit s mkpart loader2 ${LOADER2_START} $(expr ${ATF_START} - 1)
parted -s ${IMAGE} unit s mkpart atf ${ATF_START} $(expr ${BOOT_START} - 1)
parted -s ${IMAGE} unit s mkpart boot ${BOOT_START} $(expr ${ROOTFS_START} - 1)

# 将第6个分区(boot分区)标记为引导分区（boot flag）
parted -s ${IMAGE} set 6 boot on

# 创建rootfs分区，占用剩余所有空间
parted -s ${IMAGE} unit s mkpart rootfs ${ROOTFS_START} 100%

# 定义rootfs分区UUID，使用固定UUID，方便后续挂载识别
ROOT_UUID="B921B045-1DF0-41C3-AF44-4C6F280D3FAE"

# 使用gdisk的交互模式修改第7分区（rootfs分区）的UUID为ROOT_UUID并保存更改
gdisk ${IMAGE} <<EOF
x
c
7
${ROOT_UUID}
w
y
EOF


# 烧录uboot相关的镜像文件到镜像对应偏移位置
echo -e "\e[1;34m 烧录uboot相关的镜像文件到镜像对应偏移位置.....\e[0m"
dd if=${OUTPUT}/u-boot/idbloader.img of=${IMAGE} seek=${LOADER1_START} conv=notrunc
dd if=${OUTPUT}/u-boot/uboot.img of=${IMAGE} seek=${LOADER2_START} conv=notrunc
dd if=${OUTPUT}/u-boot/trust.img of=${IMAGE} seek=${ATF_START} conv=notrunc


# 创建rootfs文件系统大小，磁盘总大小减去128 MiB保留空间
echo -e "\e[1;34m 创建rootfs文件系统大小，磁盘总大小减去128 MiB保留空间.....\e[0m"
rootfs_size=$((disk_size-128))

# 创建一个大小为rootfs_size的空文件，用作rootfs分区映像文件
dd if=/dev/zero bs=1M count=$rootfs_size of=${IMAGE}2

# 格式化rootfs分区为ext4文件系统，使用4KB块大小，配置优化参数，并打上标签“rootfs”
mkfs.ext4 -F -b 4096 -E stride=2,stripe-width=1024 -L rootfs ${IMAGE}2

# 若临时挂载目录不存在，创建该目录
if [ ! -d /media/tmp ]; then
    mkdir -p /media/tmp
fi

# 挂载rootfs分区镜像到临时目录，供后续复制文件使用
mount -t ext4 ${IMAGE}2 /media/tmp

# 将rootfs目录中的所有文件复制到挂载点，递归复制并保留文件属性
cp -rfa $OUTPUT/rootfs/* /media/tmp

# 复制wifi相关系统文件到挂载点以供系统使用
cp -rfa $ROOT/external/system/ /media/tmp 

# 卸载挂载点，确保数据写回镜像文件中
umount /media/tmp

# 调用生成boot镜像脚本，传入第四个参数
./Generate_boot_image.sh $VAR

# 将生成的boot.img写回到镜像文件的boot分区起始位置，保持不截断已写数据
dd if=${OUTPUT}/boot.img of=${IMAGE} conv=notrunc seek=${BOOT_START}

# 将rootfs分区映像写回镜像文件的rootfs分区起始位置
dd if=${IMAGE}2 of=${IMAGE} seek=${ROOTFS_START}

# gdisk ${IMAGE} <<EOF
# p
# v
# w
# y
# EOF

# # 修复备份 GPT 表，避免备份表损坏警告
# if command -v sgdisk >/dev/null 2>&1; then
    # echo -e "\e[1;34m 修复备份 GPT 表.....\e[0m"
    # sgdisk -e "$IMAGE"
# else
    # echo "Warning: 'sgdisk' not found, cannot fix backup GPT header automatically."
# fi

# 进入镜像输出目录
cd $OUTPUT/images/ 

# 删除旧的压缩包文件，避免覆盖报错
rm -rf ${IMAGENAME}.tar.gz

# 计算镜像文件MD5值，生成校验文件，便于后续验证镜像完整性
md5sum ${IMAGENAME} > ${IMAGENAME}.md5sum

# 打包镜像文件及相关文件为tar.gz压缩包，方便发布和传输
tar czvf  ${IMAGENAME}.tar.gz $IMAGENAME*

# 删除MD5校验文件，保持目录清洁
rm -f ${IMAGENAME}.md5sum
rm -f ${IMAGE}2

# 同步文件系统，确保所有数据写入物理介质
sync

# 清空终端屏幕，脚本完成
clear

if [ $? -eq 0 ]; then
  echo -e "\e[1;32m 构建镜像成功，输出路径$OUTPUT/images/.....\e[0m"
else
  echo "命令执行失败，错误码：$?"
fi

# 清空终端屏幕，脚本完成
# clear